"""
Provide process management
"""

import ctypes
import sys

import win32api
import win32gui
import win32process

_pAddressGetCommandLinePlusOne	= win32api.GetProcAddress(win32gui.GetModuleHandle('kernel32'), 'GetCommandLineA') + 1

# declare win32 API 
CloseHandle				= ctypes.windll.kernel32.CloseHandle
CreateToolhelp32Snapshot= ctypes.windll.kernel32.CreateToolhelp32Snapshot
OpenProcess				= ctypes.windll.kernel32.OpenProcess
Process32First			= ctypes.windll.kernel32.Process32First
Process32Next			= ctypes.windll.kernel32.Process32Next
ReadProcessMemory		= ctypes.windll.kernel32.ReadProcessMemory
TerminateProcess		= ctypes.windll.kernel32.TerminateProcess
Thread32First			= ctypes.windll.kernel32.Thread32First	
Thread32Next			= ctypes.windll.kernel32.Thread32Next

#################################################################################################

class PROCESSENTRY32(ctypes.Structure):
	_fields_	=[("dwSize", ctypes.c_ulong),
	("cntUsage", ctypes.c_ulong),
	("th32ProcessID", ctypes.c_ulong),
	("th32DefaultHeapID", ctypes.c_ulong),
	("th32ModuleID", ctypes.c_ulong),
	("cntThreads", ctypes.c_ulong),
	("th32ParentProcessID", ctypes.c_ulong),
	("pcPriClassBase", ctypes.c_ulong),
	("dwFlags", ctypes.c_ulong),
	("szExeFile", ctypes.c_char * 260)
	]

class THREADENTRY32(ctypes.Structure):
	_fields_	=[("dwSize", ctypes.c_ulong),
	("cntUsage", ctypes.c_ulong),
	("th32ThreadID", ctypes.c_ulong),
	("th32OwnerProcessID", ctypes.c_ulong),
	("tpBasePri", ctypes.c_ulong),
	("tpDeltaPri", ctypes.c_ulong),
	("dwFlags", ctypes.c_ulong),
	]

#################################################################################################
def _get_pe32_name(pe32):
	'''
	From python3.2, pe32.szExeFile will be of type 'bytes', so we need to decode it manually
	'''
	name	= pe32.szExeFile
	return name.decode(sys.getfilesystemencoding()) if isinstance(name, bytes) else name

class process:
	def __init__(self, pe32):
		'''
		initialize the process object with a PROCESSENTRY32 struct
		'''
		global _get_pe32_name

		if not isinstance(pe32, PROCESSENTRY32):
			raise TypeError("invalid pe32 type (%s, expected PROCESSENTRY32)" % type(pe32))

		self.pid	= pe32.th32ProcessID
		self.id		= self.pid
		self.name	= _get_pe32_name(pe32)
	
	@property
	def commandline(self):
		if hasattr(self, '_commandline'):
			return self._commandline

		PROCESS_ALL_PROCESS	= 0x001F0FFF
		hProcess			= OpenProcess(PROCESS_ALL_PROCESS, False, self.id)
		if not hProcess:
			self._commandline	= ''
			return self._commandline

		pEnvironment		= ctypes.c_ulong(0)
		pBuffer				= (ctypes.c_char * 512)()
		ReadProcessMemory(hProcess, _pAddressGetCommandLinePlusOne, ctypes.byref(pEnvironment), ctypes.sizeof(pEnvironment), 0)
		ReadProcessMemory(hProcess, pEnvironment.value, ctypes.byref(pEnvironment), ctypes.sizeof(pEnvironment), 0)
		ReadProcessMemory(hProcess, pEnvironment.value, ctypes.byref(pBuffer), ctypes.sizeof(pBuffer), 0)
		encoding			= sys.getfilesystemencoding()
		self._commandline	= pBuffer.value.decode(encoding)
		return self._commandline

	def kill(self):
		'''
		kill the current process with self.pid
		'''
		return kill(self.pid)
	
	def listthreads(self):
		"""
		return a thread id list owning to the given pid 
		"""
		return listthreads(self.pid)

	def listwindows(self):
		'''
		list a list of window handles owning to the current process
		'''
		return listwindows(self.pid)

#################################################################################################

def listall(process_name= None):			
	"""
	return a process list which contains all running processes in the current operating system
	"""
	
	if process_name and not isinstance(process_name, str):
		raise TypeError("invalid process_name type (%s, expected string)" % type(process_name))

	
	global _get_pe32_name

	# declare enum variables
	TH32CS_SNAPPROCESS		= 0x00000002
	hProcessSnap			= CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)
	pe32					= PROCESSENTRY32()
	pe32.dwSize				= ctypes.sizeof(PROCESSENTRY32)

	if Process32First(hProcessSnap, ctypes.byref(pe32)):
		# filter by process_name
		if process_name:
			process_name	= process_name.lower()

			if _get_pe32_name(pe32).lower() == process_name:
				yield process(pe32)
			while Process32Next(hProcessSnap,ctypes.byref(pe32)):
				if _get_pe32_name(pe32).lower() == process_name:
					yield process(pe32)

		# return all available processes
		else:
			yield process(pe32)
			while Process32Next(hProcessSnap,ctypes.byref(pe32)):
				yield process(pe32)
	else:
		raise OSError("Failed to call Process32First.")
	
	CloseHandle(hProcessSnap)

#################################################################################################
def listthreads(pid):			 
	"""
	return a thread id list owning to the given pid 
	"""
	
	if not isinstance(pid, int):
		raise TypeError("invalid pid type (%s, expected int)" % type(pid))

	TH32CS_SNAPTHREAD		= 0x00000004
	hThreadSnap				= CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0)
	te32					= THREADENTRY32()
	te32.dwSize				= ctypes.sizeof(THREADENTRY32)

	if Thread32First(hThreadSnap, ctypes.byref(te32)):
		if te32.th32OwnerProcessID == pid:
			yield te32.th32ThreadID

		while Thread32Next(hThreadSnap,ctypes.byref(te32)):
			if te32.th32OwnerProcessID == pid:
				yield te32.th32ThreadID
	else:
		raise OSError("Failed to call Thread32First.")
	
	CloseHandle(hThreadSnap)

#################################################################################################

def listwindows(pid):			 
	"""
	return a windows handle list owning to the given pid 
	"""

	def _windowEnumerationHandler(hwnd, handles):
		idThread, idProcess = win32process.GetWindowThreadProcessId(hwnd)
		if idProcess == pid:
			handles.append(hwnd)
			#print(win32gui.GetWindowText(hwnd))

	#real process
	handles = []
	win32gui.EnumWindows(_windowEnumerationHandler, handles)

	return handles

#################################################################################################

def kill(pid):
	'''
	Kill a process by pid
	'''
	if not isinstance(pid, int):
		raise TypeError("invalid pid type (%s, expected int)" % type(pid))

	PROCESS_TERMINATE	= 1
	handle				= OpenProcess(PROCESS_TERMINATE, False, pid)
	if handle:
		ret				= TerminateProcess(handle, 1)
		isTerminated	= ret != 0
		return isTerminated

	return False

#################################################################################################
# only for test
if __name__=='__main__':
	# name	= 'cmd.exe'
	name	= None

	for p in listall(name):
		print(p.id, p.name, p.commandline)
